# coding = utf-8

'''
针对HDenseUnet的预处理
'''
import os
import numpy as np


def generate_livertxt(image_path, save_folder):
    if not os.path.exists(save_folder):
        os.mkdir(save_folder)

    # Generate Livertxt
    Livertxt = os.path.join(save_folder, 'LiverPixels')
    if not os.path.exists(Livertxt):
        os.mkdir(Livertxt)

    for item in os.listdir(image_path):
        item_path = os.path.join(image_path, item)
        liver = np.load(os.path.join(item_path,"segment.npy"))
        index = np.where((liver == 1) | (liver == 2))
        x = index[1]
        y = index[2]
        z = index[0]

        f = open(os.path.join(Livertxt,"{}.txt".format(item)), 'w')
        np.savetxt(f, np.c_[x, y, z], fmt="%d")
        f.write("\n")
        f.close()

        print("save {}".format(item))

def generate_tumortxt(image_path, save_folder):
    if not os.path.exists(save_folder):
        os.mkdir(save_folder)

    # Generate Livertxt
    Livertxt = os.path.join(save_folder, 'TumorPixels')
    if not os.path.exists(Livertxt):
        os.mkdir(Livertxt)

    for item in os.listdir(image_path):
        item_path = os.path.join(image_path, item)
        liver = np.load(os.path.join(item_path,"segment.npy"))
        index = np.where(liver == 2)
        x = index[1]
        y = index[2]
        z = index[0]

        f = open(os.path.join(Livertxt,"{}.txt".format(item)), 'w')
        np.savetxt(f, np.c_[x, y, z], fmt="%d")
        f.write("\n")
        f.close()

        print("save {}".format(item))

def generate_boxtxt(image_path, save_folder):
    if not os.path.exists(save_folder):
        os.mkdir(save_folder)

    # Generate Livertxt
    BoxTxt = os.path.join(save_folder, "LiverBox")
    if not os.path.exists(BoxTxt):
        os.mkdir(BoxTxt)
    for i in range(0,20):
        file_name = os.path.join(image_path, "case_{}.txt".format(str(i).zfill(5)))
        values = np.loadtxt(file_name, delimiter=' ', usecols=[0, 1, 2])
        a = np.min(values, axis=0)
        b = np.max(values, axis=0)
        box = np.append(a,b, axis=0)
        print(box)
        np.savetxt(os.path.join(BoxTxt, "case_{}.txt".format(str(i).zfill(5))), box,fmt='%d')

if __name__ == '__main__':
    #generate_tumortxt(image_path="/datasets/3Dircadb/origion", save_folder="/datasets/3Dircadb/HDenseNet")
    generate_boxtxt(image_path="/datasets/3Dircadb/HDenseNet/LiverPixels", save_folder="/datasets/3Dircadb/HDenseNet")
